import gradio as gr
import os
from openai import OpenAI
import json
import uvicorn
from fastapi import FastAPI, HTTPException, UploadFile, File
from pydantic import BaseModel
from typing import Optional, List
from PIL import Image
import io

# --- 配置部分 ---
# 从环境变量或直接设置您的百度AI API Key和Secret Key
BAIDU_API_KEY = os.environ.get("BAIDU_API_KEY", "xxx")  # 请替换为您的API Key
BAIDU_BASE_URL = "https://aistudio.baidu.com/llm/lmapi/v3"
MODEL_NAME = "ernie-4.5-turbo-vl"  # 百度模型名称

# 初始化 OpenAI 客户端以兼容百度API
client = OpenAI(
    api_key=BAIDU_API_KEY,
    base_url=BAIDU_BASE_URL,
)

# 支持的产品类型和常见缺陷
PRODUCT_TYPES = {
    "电子元件": ["外观损伤", "尺寸偏差", "焊接不良", "引脚变形", "标识错误"],
    "机械零件": ["表面划痕", "裂纹", "变形", "毛刺", "装配错误"],
    "塑料产品": ["气泡", "色差", "缩水", "飞边", "断裂"],
    "金属制品": ["锈蚀", "凹陷", "划痕", "镀层缺陷", "变形"],
    "纺织品": ["断线", "污渍", "色差", "针脚不齐", "瑕疵"],
    "包装产品": ["密封不良", "标签错误", "破损", "印刷错误", "尺寸不符"],
}


# --- 核心处理函数 ---
def analyze_quality_core(product_type, defects, description, image_data=None):
    """
    质检分析核心逻辑
    """
    try:
        # 构建提示词
        prompt = f"""
        你是一位资深工业质检专家，擅长分析各类产品的质量问题。
        
        请根据以下信息，对产品质量进行分析并提供专业建议：
        - 产品类型：{product_type}
        - 关注缺陷类型：{', '.join(defects) if defects else '未指定，需要全面检查'}
        - 产品描述：{description if description else '无特殊描述'}
        
        请按照以下结构提供分析结果：
        1. 潜在质量问题：列出可能存在的质量问题及风险等级（高/中/低）
        2. 检测建议：针对这些问题的具体检测方法和标准
        3. 改进建议：生产过程中可采取的改进措施
        4. 处理方案：对已发现问题的产品的处理建议
        
        分析时请结合产品类型特点，语言专业、准确，建议具有可操作性。
        """

        # 调用百度AI API
        messages = [{"role": "user", "content": prompt}]

        # 如果有图片，添加图片分析请求（简化处理）
        if image_data:
            messages[0][
                "content"
            ] += "\n\n同时请分析提供的产品图片，识别是否存在可见缺陷。"

        completion = client.chat.completions.create(
            model=MODEL_NAME,
            messages=messages,
            stream=False,
            temperature=0.4,  # 降低创造性，提高准确性
        )

        ai_response = completion.choices[0].message.content
        return {"success": True, "result": ai_response}
    except Exception as e:
        print(f"调用AI API时出错: {e}")
        return {"success": False, "error": f"分析失败：{str(e)}"}


# --- Gradio界面函数 ---
def gradio_interface():
    with gr.Blocks(title="质检慧眼 - 工业质检AI助手") as demo:
        gr.Markdown("# 🔍 质检慧眼 - 工业质检AI助手")
        gr.Markdown("上传产品图片或输入产品信息，获取专业的质量检测分析和建议")

        with gr.Row():
            with gr.Column(scale=1):
                product_type = gr.Dropdown(
                    label="产品类型",
                    choices=list(PRODUCT_TYPES.keys()),
                    value="电子元件",
                )

                defects = gr.CheckboxGroup(
                    label="关注缺陷类型",
                    choices=PRODUCT_TYPES["电子元件"],
                    value=PRODUCT_TYPES["电子元件"][:2],
                )

                description = gr.Textbox(
                    label="产品描述（可选）",
                    placeholder="请描述产品的基本情况、生产工艺等...",
                )

                image = gr.Image(label="产品图片（可选）", type="pil", height=200)

                analyze_btn = gr.Button("开始分析", variant="primary")

            with gr.Column(scale=2):
                result = gr.Markdown(label="质检分析结果")

        # 当产品类型变化时，更新缺陷选项
        def update_defects(product_type):
            return gr.CheckboxGroup(
                choices=PRODUCT_TYPES[product_type],
                value=PRODUCT_TYPES[product_type][:2],
            )

        product_type.change(fn=update_defects, inputs=[product_type], outputs=[defects])

        # 分析按钮点击事件
        def analyze(product_type, defects, description, image):
            return analyze_quality_core(product_type, defects, description, image)[
                "result"
            ]

        analyze_btn.click(
            fn=analyze,
            inputs=[product_type, defects, description, image],
            outputs=[result],
        )

        # API说明
        gr.Markdown(
            """
        ---
        ## API 接口说明
        可通过API接口集成到生产系统：
        - **POST** `/api/analyze`
        - **参数**: 产品类型、缺陷类型、描述、图片(可选)
        - **返回**: 质检分析结果
        """
        )

    return demo


# --- FastAPI 应用 ---
app = FastAPI(title="质检慧眼 - 工业质检AI服务", description="提供工业产品质量检测分析")


# 请求模型
class AnalysisRequest(BaseModel):
    product_type: str
    defects: Optional[List[str]] = []
    description: Optional[str] = ""


# 分析接口
@app.post("/api/analyze", summary="产品质量分析")
async def api_analyze(
    request: AnalysisRequest, file: Optional[UploadFile] = File(None)
):
    image_data = None
    if file:
        # 处理上传的图片
        contents = await file.read()
        image_data = Image.open(io.BytesIO(contents))

    result = analyze_quality_core(
        request.product_type, request.defects, request.description, image_data
    )

    if result["success"]:
        return result
    else:
        raise HTTPException(status_code=500, detail=result["error"])


# 健康检查
@app.get("/health", summary="健康检查")
async def health_check():
    return {"status": "healthy"}


# 挂载Gradio界面
demo = gradio_interface()
app = gr.mount_gradio_app(app=app, blocks=demo, path="/")

# 启动应用
if __name__ == "__main__":
    print("🚀 启动质检慧眼服务...")
    print("📱 Gradio 界面: http://localhost:7860")
    print("💻 API 文档: http://localhost:7860/docs")
    uvicorn.run(app, host="0.0.0.0", port=7860)
